#include "KeyFrame.h"
#include "Converter.h"



using namespace std;
using namespace cv;


static long unsigned int  g_KeyframeIdCount = 0;


KeyFrame::KeyFrame(Frame &F, Map *pMap):
	m_frameId(F.m_frameId),
	m_HalfBaseline(F.m_baseline/2),
	m_keyPointsNumber(F.m_keyPointsNumber),
	m_fx(F.m_fx),
	m_fy(F.m_fy),
	m_cx(F.m_cx),
	m_cy(F.m_cy),
	m_invfx(F.m_invfx),
	m_invfy(F.m_invfy),
	m_baselineFx(F.m_baselineFx),
	m_baseline(F.m_baseline),
	m_Map(pMap),
	m_vectorKeyPoints(F.m_vectorKeyPointsLeft),
	m_vectorKeyPointsUn(F.m_vectorKeyPointsLeftUn),
	m_vectorURight(F.m_vectorURight),
	m_vectorDepth(F.m_vectorDepth),
	m_Descriptors(F.m_DescriptorsLeft.clone()),
	m_K(F.m_K),
	m_Parent(NULL),
	m_FirstConnection(true),
	m_NotErase(false),
	m_ToBeErased(false),
	m_BadKeyFrame(false),
	m_vectorMapPoints(F.m_vectorMapPoints),
	m_TrackReferenceForFrame(0),
	m_FuseTargetForKF(0),
	m_ThDepth(F.m_ThDepth),
	m_BALocalForKF(0),
	m_BAFixedForKF(0),
	m_MaxX(F.m_MaxX),
	m_MaxY(F.m_MaxY),
	m_MinX(F.m_MinX),
	m_MinY(F.m_MinY),
	m_GridCols(FRAME_GRID_COLS),
	m_GridRows(FRAME_GRID_ROWS),
	m_GridElementWidthInv(F.m_GridElementWidthInv),
	m_GridElementHeightInv(F.m_GridElementHeightInv)
{
	m_KeyFrameId = g_KeyframeIdCount++;

	m_Grid.resize(m_GridCols);
    for(int i=0; i<m_GridCols;i++)
    {
        m_Grid[i].resize(m_GridRows);
        for(int j=0; j<m_GridRows; j++)
            m_Grid[i][j] = F.m_Grid[i][j];
    }

	SetPose(F.m_RTcw);
}


KeyFrame::~KeyFrame()
{


}

cv::Mat KeyFrame::GetCameraCenter()
{
	unique_lock<mutex> lock(m_MutexPose);
	return m_Ow.clone();
}


void KeyFrame::SetPose(const cv::Mat &RTcw)
{
	 unique_lock<mutex> lock(m_MutexPose);
	 RTcw.copyTo(m_RTcw);
	 cv::Mat Rcw = m_RTcw.rowRange(0,3).colRange(0,3);
	 cv::Mat Tcw = m_RTcw.rowRange(0,3).col(3);
	 cv::Mat Rwc = Rcw.t();
	 m_Ow = -Rwc * Tcw;

	 m_RTwc = cv::Mat::eye(4,4,RTcw.type());
	 Rwc.copyTo(m_RTwc.rowRange(0,3).colRange(0,3));
	 m_Ow.copyTo(m_RTwc.rowRange(0,3).col(3));
	 cv::Mat center = (cv::Mat_<float>(4,1) << m_HalfBaseline, 0 , 0, 1);
	 m_Cw = m_RTwc*center;

	 m_position.x = Tcw.at<float>(0);
	 m_position.y = Tcw.at<float>(1);
	 m_position.z = Tcw.at<float>(2);

	 vector<float> q = Converter::toQuaternion(Rwc);

	Converter::getRollPitchYawFromQuaternion(q[0], q[1], q[2], q[3], m_attitude.x, m_attitude.y, m_attitude.z);
}


cv::Mat KeyFrame::GetPose()
{
	unique_lock<mutex> lock(m_MutexPose);
	return m_RTcw.clone();
}


cv::Mat KeyFrame::GetPoseInverse()
{
	 unique_lock<mutex> lock(m_MutexPose);
	 return m_RTwc.clone();
}


cv::Mat KeyFrame::GetStereoCenter()
{
	unique_lock<mutex> lock(m_MutexPose);
	return m_Cw.clone();
}


cv::Mat KeyFrame::GetRotation()
{
	unique_lock<mutex> lock(m_MutexPose);
	return m_RTcw.rowRange(0,3).colRange(0,3).clone();
}


cv::Mat KeyFrame::GetTranslation()
{
	unique_lock<mutex> lock(m_MutexPose);
	return m_RTcw.rowRange(0,3).col(3).clone();
}


void KeyFrame::AddConnection(KeyFrame* pKF, const int &weight)
{
	{
		unique_lock<mutex> lock(m_MutexConnections);
		if(!m_ConnectedKeyFrameWeights.count(pKF))
		{
			m_ConnectedKeyFrameWeights[pKF] = weight;
		}
		else if(m_ConnectedKeyFrameWeights[pKF] != weight)
		{
			m_ConnectedKeyFrameWeights[pKF] = weight;
		}
		else
		{
			return;
		}
	}

	UpdateBestCovisibles();
}


void KeyFrame::EraseConnection(KeyFrame* pKF)
{
	bool bUpdate = false;

	{
		unique_lock<mutex> lock(m_MutexConnections);
		if(m_ConnectedKeyFrameWeights.count(pKF))
		{
			m_ConnectedKeyFrameWeights.erase(pKF);
			bUpdate = true;
		}
	}

	if(bUpdate)
	{
		UpdateBestCovisibles();
	}
}


void KeyFrame::UpdateConnections()
{
	map<KeyFrame*,int> KFcounter;
	vector<MapPoint*> vpMP;

	{
		unique_lock<mutex> lockMPs(m_MutexFeatures);
		vpMP = m_vectorMapPoints;
	}

	//For all map points in keyframe check in which other keyframes are they seen
	//Increase counter for those keyframes
	for(vector<MapPoint*>::iterator it=vpMP.begin(), end=vpMP.end(); it!=end; it++)
	{
		MapPoint* pMP = *it;

		if(!pMP)
		{
			continue;
		}

		if(pMP->isBadMapPoint())
		{
			continue;
		}

		map<KeyFrame*,size_t> observations = pMP->GetObservations();
		for(map<KeyFrame*,size_t>::iterator mit=observations.begin(), mend=observations.end(); mit!=mend; mit++)
		{
			if(mit->first->m_KeyFrameId == m_KeyFrameId)
			{
				continue;
			}

			KFcounter[mit->first]++;
		}
	}

	if(KFcounter.empty())
	{
		return;
	}

	//If the counter is greater than threshold add connection
	//In case no keyframe counter is over threshold add the one with maximum counter
	int nmax=0;
	KeyFrame* pKFmax=NULL;
	int th = 15;

	vector<pair<int,KeyFrame*> > vPairs;
	vPairs.reserve(KFcounter.size());
	for(map<KeyFrame*,int>::iterator it=KFcounter.begin(), end=KFcounter.end(); it!=end; it++)
	{
		if(it->second > nmax)
		{
			nmax = it->second;
			pKFmax = it->first;
		}

		if(it->second >= th)
		{
			vPairs.push_back(make_pair(it->second, it->first));
			(it->first)->AddConnection(this,it->second);
		}
	}

	if(vPairs.empty())
	{
		vPairs.push_back(make_pair(nmax,pKFmax));
		pKFmax->AddConnection(this,nmax);
	}

	sort(vPairs.begin(),vPairs.end());
	list<KeyFrame*> KFs;
	list<int> Ws;
	for(size_t i=0; i<vPairs.size();i++)
	{
		KFs.push_front(vPairs[i].second);
		Ws.push_front(vPairs[i].first);
	}

	{
		unique_lock<mutex> lockCon(m_MutexConnections);
		m_ConnectedKeyFrameWeights = KFcounter;
		m_vectorOrderedConnectedKeyFrames = vector<KeyFrame*>(KFs.begin(),KFs.end());
		m_vectorOrderedWeights = vector<int>(Ws.begin(), Ws.end());

		if(m_FirstConnection && m_KeyFrameId!=0)
		{
			m_Parent = m_vectorOrderedConnectedKeyFrames.front();
			m_Parent->AddChild(this);
			m_FirstConnection = false;
		}
	}
}


void KeyFrame::UpdateBestCovisibles()
{
	unique_lock<mutex> lock(m_MutexConnections);

	vector<pair<int,KeyFrame*> > vPairs;
	vPairs.reserve(m_ConnectedKeyFrameWeights.size());

	for(map<KeyFrame*,int>::iterator it=m_ConnectedKeyFrameWeights.begin(), end=m_ConnectedKeyFrameWeights.end(); it!=end; it++)
	{
		vPairs.push_back(make_pair(it->second,it->first));
	}

	sort(vPairs.begin(),vPairs.end());

	list<KeyFrame*> KFs;
	list<int> Ws;
	for(size_t i=0; i<vPairs.size(); i++)
	{
		KFs.push_front(vPairs[i].second);
		Ws.push_front(vPairs[i].first);
	}

	m_vectorOrderedConnectedKeyFrames = vector<KeyFrame*>(KFs.begin(),KFs.end());
	m_vectorOrderedWeights = vector<int>(Ws.begin(), Ws.end()); 
}

std::set<KeyFrame *> KeyFrame::GetConnectedKeyFrames()
{
	unique_lock<mutex> lock(m_MutexConnections);
	set<KeyFrame*> s;
	for(map<KeyFrame*,int>::iterator it=m_ConnectedKeyFrameWeights.begin(); it!=m_ConnectedKeyFrameWeights.end(); it++)
	{
		s.insert(it->first);
	}

	return s;
}

std::vector<KeyFrame* > KeyFrame::GetVectorCovisibleKeyFrames()
{
	unique_lock<mutex> lock(m_MutexConnections);
	return m_vectorOrderedConnectedKeyFrames;
}


std::vector<KeyFrame*> KeyFrame::GetBestCovisibilityKeyFrames(const int &N)
{
	unique_lock<mutex> lock(m_MutexConnections);
	if((int)m_vectorOrderedConnectedKeyFrames.size() < N)
	{
		return m_vectorOrderedConnectedKeyFrames;
	}
	else
	{
		return vector<KeyFrame*>(m_vectorOrderedConnectedKeyFrames.begin(), m_vectorOrderedConnectedKeyFrames.begin()+N);
	}
}


std::vector<KeyFrame*> KeyFrame::GetCovisiblesByWeight(const int &w)
{
	unique_lock<mutex> lock(m_MutexConnections);
	if(m_vectorOrderedConnectedKeyFrames.empty())
	{
		return vector<KeyFrame*>();
	}

	vector<int>::iterator it = upper_bound(m_vectorOrderedWeights.begin(), m_vectorOrderedWeights.end(), w, KeyFrame::weightComp);
	if(it == m_vectorOrderedWeights.end())
	{
		return vector<KeyFrame*>();
	}
	else
	{
		int n = it - m_vectorOrderedWeights.begin();
		return vector<KeyFrame*>(m_vectorOrderedConnectedKeyFrames.begin(), m_vectorOrderedConnectedKeyFrames.begin() + n);
	}
}


int KeyFrame::GetWeight(KeyFrame* pKF)
{
	unique_lock<mutex> lock(m_MutexConnections);
	if(m_ConnectedKeyFrameWeights.count(pKF))
	{
		return m_ConnectedKeyFrameWeights[pKF];
	}
	else
	{
		return 0;
	}
}


void KeyFrame::AddChild(KeyFrame* pKF)
{
	unique_lock<mutex> lockCon(m_MutexConnections);
	m_setChildrens.insert(pKF);
}


void KeyFrame::EraseChild(KeyFrame* pKF)
{
	unique_lock<mutex> lockCon(m_MutexConnections);
	m_setChildrens.erase(pKF);
}


void KeyFrame::ChangeParent(KeyFrame* pKF)
{
	unique_lock<mutex> lockCon(m_MutexConnections);
	m_Parent = pKF;
	pKF->AddChild(this);
}


std::set<KeyFrame*> KeyFrame::GetChilds()
{
	unique_lock<mutex> lockCon(m_MutexConnections);
	return m_setChildrens;
}


KeyFrame* KeyFrame::GetParent()
{
	unique_lock<mutex> lockCon(m_MutexConnections);
	return m_Parent;
}


bool KeyFrame::hasChild(KeyFrame* pKF)
{
	unique_lock<mutex> lockCon(m_MutexConnections);
	return m_setChildrens.count(pKF);
}


void KeyFrame::AddLoopEdge(KeyFrame* pKF)
{
	unique_lock<mutex> lockCon(m_MutexConnections);
	m_NotErase = true;
	m_setLoopEdges.insert(pKF);
}


std::set<KeyFrame*> KeyFrame::GetLoopEdges()
{
	unique_lock<mutex> lockCon(m_MutexConnections);
	return m_setLoopEdges;
}


void KeyFrame::AddMapPoint(MapPoint* pMP, const size_t &idx)
{
	unique_lock<mutex> lock(m_MutexFeatures);
	m_vectorMapPoints[idx] = pMP;
}


void KeyFrame::EraseMapPointMatch(const size_t &idx)
{
	unique_lock<mutex> lock(m_MutexFeatures);
	m_vectorMapPoints[idx] = static_cast<MapPoint*>(NULL);
}

void KeyFrame::EraseMapPointMatch(MapPoint* pMP)
{
	int idx = pMP->GetIndexInKeyFrame(this);
	if(idx>=0)
	{
		m_vectorMapPoints[idx] = static_cast<MapPoint*>(NULL);
	}
}

void KeyFrame::ReplaceMapPointMatch(const size_t &idx, MapPoint* pMP)
{
	m_vectorMapPoints[idx] = pMP;
}


std::set<MapPoint*> KeyFrame::GetMapPoints()
{
	unique_lock<mutex> lock(m_MutexFeatures);
	set<MapPoint*> s;
	for(size_t i=0; i<m_vectorMapPoints.size(); i++)
	{
		if(!m_vectorMapPoints[i])
		{
			continue;
		}

		MapPoint* pMP = m_vectorMapPoints[i];
		if(!pMP->isBadMapPoint())
		{
			s.insert(pMP);
		}
	}

	return s;
}


std::vector<MapPoint*> KeyFrame::GetMapPointMatches()
{
	unique_lock<mutex> lock(m_MutexFeatures);
	return m_vectorMapPoints;
}


int KeyFrame::TrackedMapPoints(const int &minObs)
{
	unique_lock<mutex> lock(m_MutexFeatures);

	int nPoints=0;
	const bool bCheckObs = minObs > 0;
	for(int i=0; i<m_keyPointsNumber; i++)
	{
		MapPoint* pMP = m_vectorMapPoints[i];
		if(pMP)
		{
			if(!pMP->isBadMapPoint())
			{
				if(bCheckObs)
				{
					if(m_vectorMapPoints[i]->Observations() >= minObs)
					{
						nPoints++;
					}
				}
				else
				{
					nPoints++;
				}
			}
		}
	}

	return nPoints;
}


MapPoint* KeyFrame::GetMapPoint(const size_t &idx)
{
	unique_lock<mutex> lock(m_MutexFeatures);
	return m_vectorMapPoints[idx];
}


cv::Mat KeyFrame::UnprojectStereo(int i)
{
	const float z = m_vectorDepth[i];
	if(z>0)
	{
		float u = m_vectorKeyPoints[i].pt.x;
		float v = m_vectorKeyPoints[i].pt.y;
		float x = (u - m_cx) * z * m_invfx;
		float y = (v - m_cy) * z * m_invfy;
        cv::Mat x3Dc = (cv::Mat_<float>(3,1) << x, y, z);

		unique_lock<mutex> lock(m_MutexPose);
		return m_RTwc.rowRange(0,3).colRange(0,3) * x3Dc + m_RTwc.rowRange(0,3).col(3);
	}
	else
	{
		return cv::Mat();
	}
}

void KeyFrame::SetNotErase()
{
	unique_lock<mutex> lock(m_MutexConnections);
	m_NotErase = true;
}

void KeyFrame::SetErase()
{
	{
		unique_lock<mutex> lock(m_MutexConnections);
		if(m_setLoopEdges.empty())
		{
			m_NotErase = false;
		}
	}

	if(m_ToBeErased)
	{
		SetBadFlag();
	}
}


void KeyFrame::SetBadFlag()
{
	{
		unique_lock<mutex> lock(m_MutexConnections);
		if(m_KeyFrameId == 0)
		{
			return;
		}
		else if(m_NotErase)
		{
			m_ToBeErased = true;
			return;
		}
	}

	for(map<KeyFrame*,int>::iterator it = m_ConnectedKeyFrameWeights.begin(), end=m_ConnectedKeyFrameWeights.end(); it!=end; it++)
	{
		it->first->EraseConnection(this);
	}

	for(size_t i=0; i<m_vectorMapPoints.size(); i++)
	{
		if(m_vectorMapPoints[i])
		{
			m_vectorMapPoints[i]->EraseObservation(this);
		}
	}

	{
		unique_lock<mutex> lock(m_MutexConnections);
		unique_lock<mutex> lock1(m_MutexFeatures);

		m_ConnectedKeyFrameWeights.clear();
		m_vectorOrderedConnectedKeyFrames.clear();

		// Update Spanning Tree
		set<KeyFrame*> sParentCandidates;
		sParentCandidates.insert(m_Parent);

		while(!m_setChildrens.empty())
		{
			bool bContinue = false;
			int max = -1;
			KeyFrame* pC;
			KeyFrame* pP;

			for(set<KeyFrame*>::iterator it=m_setChildrens.begin(), end=m_setChildrens.end(); it!=end; it++)
			{
				KeyFrame* pKF = *it;
				if(pKF->isBadKeyFrame())
				{
					continue;
				}

				// Check if a parent candidate is connected to the keyframe
				vector<KeyFrame*> vpConnected = pKF->GetVectorCovisibleKeyFrames();
				for(size_t i=0; i<vpConnected.size(); i++)
				{
					for(set<KeyFrame*>::iterator sit=sParentCandidates.begin(), send=sParentCandidates.end(); sit!=send; sit++)
					{
						if(vpConnected[i]->m_KeyFrameId == (*sit)->m_KeyFrameId)
						{
							int w = pKF->GetWeight(vpConnected[i]);
							if(w>max)
							{
								pC = pKF;
								pP = vpConnected[i];
								max = w;
								bContinue = true;
							}
						}
					}
				}
			}

			if(bContinue)
			{
				pC->ChangeParent(pP);
				sParentCandidates.insert(pC);
				m_setChildrens.erase(pC);
			}
			else
			{
				break;
			}
		}

		// If a children has no covisibility links with any parent candidate, assign to the original parent of this KF
		if(!m_setChildrens.empty())
		{
			for(set<KeyFrame*>::iterator it=m_setChildrens.begin(); it!=m_setChildrens.end(); it++)
			{
				(*it)->ChangeParent(m_Parent);
			}
		}

		m_Parent->EraseChild(this);
		m_RTcp = m_RTcw * m_Parent->GetPoseInverse();

		m_BadKeyFrame = true;
	}

	m_Map->EraseKeyFrame(this);
}



bool KeyFrame::isBadKeyFrame()
{
	unique_lock<mutex> lock(m_MutexConnections);
	return m_BadKeyFrame;
}


bool KeyFrame::IsInImage(const float &x, const float &y)
{
	return (x>=m_MinX && x<m_MaxX && y>=m_MinY && y<m_MaxY);
}


std::vector<size_t> KeyFrame::GetFeaturesInArea(const float &x, const float  &y, const float  &r)
{
	vector<size_t> vIndices;
	vIndices.reserve(m_keyPointsNumber);

	int nMinCellX = max(0, (int)floor((x - m_MinX - r) * m_GridElementWidthInv));
	if(nMinCellX >= m_GridCols)
	{
		return vIndices;
	}

	int nMaxCellX = min((int)m_GridCols - 1, (int)ceil((x - m_MinX + r) * m_GridElementWidthInv));
	if(nMaxCellX < 0)
	{
		return vIndices;
	}

	int nMinCellY = max(0, (int)floor((y - m_MinY - r) * m_GridElementHeightInv));
	if(nMinCellY >= m_GridRows)
	{
		return vIndices;
	}

	int nMaxCellY = min((int)m_GridRows - 1, (int)ceil((y - m_MinY + r) * m_GridElementHeightInv));
	if(nMaxCellY < 0)
	{
		return vIndices;
	}

	for(int ix = nMinCellX; ix<=nMaxCellX; ix++)
	{
		for(int iy = nMinCellY; iy<=nMaxCellY; iy++)
		{
			vector<size_t> vCell = m_Grid[ix][iy];
			for(size_t j=0, jend=vCell.size(); j<jend; j++)
			{
				const cv::KeyPoint &kpUn = m_vectorKeyPointsUn[vCell[j]];
				float distx = kpUn.pt.x - x;
				float disty = kpUn.pt.y - y;

				if(fabs(distx) < r && fabs(disty) < r)
				{
					vIndices.push_back(vCell[j]);
				}
			}
		}
	}

	return vIndices;
}

